import os
os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'

import torch
import torch.nn as nn
import torch.optim as optim
import h5py
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm

class VoltageProcessor:
    @staticmethod
    def calculate_calibration_matrix(Vcav_I, Vcav_Q, Vfor_I, Vfor_Q, Vrefl_I, Vrefl_Q):
        """计算校准矩阵和校准后的电压值"""
        A_matrix = np.column_stack((
            np.append(Vfor_I, Vfor_Q),
            np.append(-Vfor_Q, Vfor_I),
            np.append(Vrefl_I, Vrefl_Q),
            np.append(-Vrefl_Q, Vrefl_I)
        ))
        y_array = np.append(Vcav_I, Vcav_Q)

        try:
            C_array = np.matmul(np.linalg.pinv(A_matrix, rcond=1e-5), y_array.transpose())
            return C_array
        except np.linalg.LinAlgError as e:
            print(f"SVD计算失败: {e}")
            return None

class WaveformDataset(Dataset):
    def __init__(self, file_path, sample_ratio=0.2):
        self.file_path = file_path
        
        with h5py.File(file_path, 'r') as f:
            self.n_sets = len([key for key in f.keys() if key.startswith('set_')])
            n_samples = int(self.n_sets * sample_ratio)
            
            np.random.seed(42)
            selected_indices = np.random.choice(self.n_sets, n_samples, replace=False)
            selected_indices.sort()
            
            self.features = []
            self.labels = []
            
            for idx in tqdm(selected_indices, desc="加载数据"):
                set_name = f'set_{idx}'
                dataset = f[set_name]
                
                # 读取波形数据
                Vc = dataset['Vc_real'][()] + 1j * dataset['Vc_imag'][()]
                Vf = dataset['Vf_star_real'][()] + 1j * dataset['Vf_star_imag'][()]
                Vr = dataset['Vr_star_real'][()] + 1j * dataset['Vr_star_imag'][()]
                
                # 构建特征向量
                features = np.concatenate([
                    np.abs(Vc), np.angle(Vc),
                    np.abs(Vf), np.angle(Vf),
                    np.abs(Vr), np.angle(Vr)
                ])
                
                # 计算X Y值作为标签
                C_array = VoltageProcessor.calculate_calibration_matrix(
                    Vc.real, Vc.imag,
                    Vf.real, Vf.imag,
                    Vr.real, Vr.imag
                )
                
                if C_array is not None:
                    X = complex(C_array[0], C_array[1])
                    Y = complex(C_array[2], C_array[3])
                else:
                    X = 0+0j
                    Y = 0+0j
                
                # 构建标签向量
                labels = np.array([X.real, X.imag, Y.real, Y.imag])
                
                self.features.append(features)
                self.labels.append(labels)
            
            self.features = torch.FloatTensor(np.array(self.features))
            self.labels = torch.FloatTensor(np.array(self.labels))
    
    def __len__(self):
        return len(self.features)
    
    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]

class XYNeuralNetwork(nn.Module):
    def __init__(self, input_size):
        super(XYNeuralNetwork, self).__init__()
        
        self.network = nn.Sequential(
            nn.Linear(input_size, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 4)  # 直接输出X和Y的实部虚部
        )
        
    def forward(self, x):
        return self.network(x)

def train_model(model, train_loader, optimizer, device, num_epochs=100):
    """训练模型"""
    criterion = nn.MSELoss()
    model.train()
    
    for epoch in range(num_epochs):
        total_loss = 0
        for features, labels in train_loader:
            features = features.to(device)
            labels = labels.to(device)
            
            outputs = model(features)
            loss = criterion(outputs, labels)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        if (epoch + 1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {total_loss/len(train_loader):.4f}')

def main():
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    # 加载数据
    data_path = r'E:\MLdata\SimData\coupling_simulation_data.h5'
    dataset = WaveformDataset(data_path)
    
    # 分割训练集和测试集
    train_size = int(0.8 * len(dataset))
    test_size = len(dataset) - train_size
    train_dataset, test_dataset = torch.utils.data.random_split(
        dataset, [train_size, test_size]
    )
    
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
    
    # 创建模型
    input_size = dataset.features.shape[1]
    model = XYNeuralNetwork(input_size).to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    # 训练模型
    print("开始训练...")
    train_model(model, train_loader, optimizer, device)
    
    # 保存模型
    save_dir = os.path.dirname(data_path)
    model_path = os.path.join(save_dir, 'xy_model.pth')
    torch.save(model.state_dict(), model_path)
    print(f"模型已保存到: {model_path}")
    
    # 评估模型并保存预测结果
    model.eval()
    predictions = []
    actuals = []
    
    with torch.no_grad():
        for features, labels in test_loader:
            features = features.to(device)
            outputs = model(features)
            predictions.append(outputs.cpu().numpy())
            actuals.append(labels.cpu().numpy())
    
    predictions = np.vstack(predictions)
    actuals = np.vstack(actuals)
    
    # 创建结果DataFrame
    results_df = pd.DataFrame({
        'X_true_real': actuals[:, 0],
        'X_true_imag': actuals[:, 1],
        'X_pred_real': predictions[:, 0],
        'X_pred_imag': predictions[:, 1],
        'Y_true_real': actuals[:, 2],
        'Y_true_imag': actuals[:, 3],
        'Y_pred_real': predictions[:, 2],
        'Y_pred_imag': predictions[:, 3]
    })
    
    # 计算相对误差
    for param in ['X', 'Y']:
        true_complex = results_df[f'{param}_true_real'] + 1j * results_df[f'{param}_true_imag']
        pred_complex = results_df[f'{param}_pred_real'] + 1j * results_df[f'{param}_pred_imag']
        relative_error = np.abs(pred_complex - true_complex) / np.abs(true_complex) * 100
        results_df[f'{param}_relative_error_%'] = relative_error
    
    # 保存结果
    save_path = os.path.join(save_dir, 'xy_predictions.csv')
    results_df.to_csv(save_path, index=False)
    print(f"\n预测结果已保存到: {save_path}")
    
    # 打印统计信息
    print("\n预测误差统计:")
    for param in ['X', 'Y']:
        errors = results_df[f'{param}_relative_error_%']
        print(f"\n{param}参数:")
        print(f"平均相对误差: {errors.mean():.2f}%")
        print(f"最大相对误差: {errors.max():.2f}%")
        print(f"中位相对误差: {errors.median():.2f}%")

if __name__ == '__main__':
    main() 